# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/23 16:47 
# @Author  : shutao
# @FileName: fsd_dataset.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import numpy as np
import cv2 as cv
import random
import json
import os
from loguru import logger

from .base_dataset import BaseDataset


class FSDDataset(BaseDataset):
    """ Action recognition """
    def __init__(self, root_path, is_test=False, test_data_npy_path=None):
        super(FSDDataset, self).__init__()
        # 设置数据集外部参数
        self.root_path = root_path
        self.is_test = is_test
        # 拼出数据文件绝对路径
        data_npy_path = os.path.join(root_path, "train_data.npy")
        label_npy_path = os.path.join(root_path, "train_label.npy")
        # 验证数据.npy文件是否存在
        assert os.path.exists(data_npy_path), "The target file does not exist"
        assert os.path.exists(label_npy_path), "The target file does not exist"
        if test_data_npy_path is not None:
            assert os.path.exists(test_data_npy_path), "The target file does not exist"

        # 根据训练程序修改数据读取内容
        if not is_test:
            self.datalist = np.load(data_npy_path)
            self.labellist = np.load(label_npy_path)
        else:
            if test_data_npy_path is not None:
                self.datalist = np.load(test_data_npy_path)
        pass

    def _process(self, data, label):
        return data, label

    def __getitem__(self, item):
        data = self.datalist[item]

        if not self.is_test:  # 训练、验证阶段
            label = self.labellist[item]

            data, label = self._process(data, label)
            return data, label
        else:  # 测试阶段
            return data

    def __len__(self):
        return self.datalist.shape[0]
